/*
  S.M.A.C.K - An operating system kernel
  Copyright (C) 2010,2011 Mattias Holm and Kristian Rietveld
  For licensing and a full list of authors of the kernel, see the files
  COPYING and AUTHORS.
*/

#include "wfs.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>


/*
 * Utilities
 */
static size_t
wfs_get_size(void)
{
  return sizeof(wfs_file_entry_t) * WFS_N_FILES
      + WFS_N_BLOCKS * sizeof(uint16_t)
      + WFS_N_BLOCKS * 512;
}

static int
wfs_file_entry_is_empty(const wfs_file_entry_t *entry)
{
  if (entry->filename[0] == 0 || entry->start_block == WFS_BLOCK_FREE)
    return 1;

  return 0;
}

static int
wfs_file_entry_is_directory(const wfs_file_entry_t *entry)
{
  if ((entry->size & WFS_SIZE_IS_DIRECTORY) == WFS_SIZE_IS_DIRECTORY)
    return 1;

  return 0;
}

static uint32_t
wfs_file_entry_get_size(const wfs_file_entry_t *entry)
{
  return entry->size & WFS_SIZE_MASK;
}

static int
wfs_check_image(const char *filename)
{
  struct stat buf;

  if (stat(filename, &buf) < 0)
    {
      fprintf(stderr, "error: file '%s': %s\n", filename, strerror(errno));
      return -1;
    }

  if (buf.st_size != wfs_get_size())
    {
      fprintf(stderr, "error: file '%s' has incorrect size\n", filename);
      return -1;
    }

  return 0;
}

static int
wfs_open_image(const char *filename)
{
  int fd;

  if (wfs_check_image(filename) < 0)
    return -1;

  fd = open(filename, O_RDWR);
  if (fd < 0)
    {
      fprintf(stderr, "error: could not open file '%s': %s\n",
              filename, strerror(errno));
      return -1;
    }

  return fd;
}

static int
wfs_entry_exists(const int fd, const char *filename)
{
  int i;

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t entry;

      read(fd, &entry, sizeof(wfs_file_entry_t));

      if (wfs_file_entry_is_empty(&entry))
        continue;

      if (strncmp(entry.filename, filename, 58) == 0)
        return 1;
    }

  return 0;
}

static int
wfs_add_entry(const int fd, const wfs_file_entry_t *entry)
{
  int i;

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t tmp_entry;

      read(fd, &tmp_entry, sizeof(wfs_file_entry_t));
      if (wfs_file_entry_is_empty(&tmp_entry))
        {
          off_t offset = -1 * (off_t)sizeof(wfs_file_entry_t);
          lseek(fd, offset, SEEK_CUR);
          write(fd, entry, sizeof(wfs_file_entry_t));
          return 0;
        }
    }

  fprintf(stderr, "error: could not allocate entry: file table is full\n");
  return -1;
}

static int
wfs_clear_entry(const int fd, const wfs_file_entry_t *entry)
{
  int i;

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t tmp_entry;

      read(fd, &tmp_entry, sizeof(wfs_file_entry_t));
      if (wfs_file_entry_is_empty(&tmp_entry))
        continue;

      if (strncmp(tmp_entry.filename, entry->filename, 58) == 0)
        {
          wfs_file_entry_t empty_entry;
          off_t offset = -1 * (off_t)sizeof(wfs_file_entry_t);

          memset(&empty_entry, 0, sizeof(wfs_file_entry_t));

          lseek(fd, offset, SEEK_CUR);
          write(fd, &empty_entry, sizeof(wfs_file_entry_t));

          return 0;
        }
    }

  return -1;
}

static int
wfs_find_entry(const int fd, const char *filename,
               wfs_file_entry_t *entry)
{
  int i;

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t tmp_entry;

      read(fd, &tmp_entry, sizeof(wfs_file_entry_t));

      if (wfs_file_entry_is_empty(&tmp_entry))
        continue;

      if (strncmp(tmp_entry.filename, filename, 58) == 0)
        {
          *entry = tmp_entry;
          return 1;
        }
    }

  return 0;
}

static int
wfs_add_or_find_entry(const int fd, wfs_file_entry_t *entry)
{
  int i;
  int free_entry = -1;

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t tmp_entry;

      read(fd, &tmp_entry, sizeof(wfs_file_entry_t));
      if (wfs_file_entry_is_empty(&tmp_entry))
        {
          if (free_entry == -1)
            free_entry = i;
          continue;
        }

      if (strncmp(tmp_entry.filename, entry->filename, 58) == 0)
        {
          off_t offset = -1 * (off_t)sizeof(wfs_file_entry_t);

          /* Get start block from existing entry */
          entry->start_block = tmp_entry.start_block;

          /* Update size in existing entry */
          tmp_entry.size = entry->size;
          lseek(fd, offset, SEEK_CUR);
          write(fd, &tmp_entry, sizeof(wfs_file_entry_t));

          return 0;
        }
    }

  if (free_entry == -1)
    {
      fprintf(stderr, "error: could not allocate entry: file table is full\n");
      return -1;
    }

  lseek(fd, free_entry * sizeof(wfs_file_entry_t), SEEK_SET);
  write(fd, entry, sizeof(wfs_file_entry_t));
  return 0;
}

static int
wfs_get_free_block(uint16_t *block_table)
{
  int i;

  for (i = 0; i < WFS_N_BLOCKS; i++)
    {
      if (block_table[i] == WFS_BLOCK_FREE)
        return i;
    }

  return -1;
}

static uint16_t *
wfs_load_block_table(int fd)
{
  uint16_t *block_table;

  block_table = malloc(WFS_BLOCK_TABLE_SIZE);
  lseek(fd, WFS_BLOCK_TABLE_START, SEEK_SET);
  read(fd, block_table, WFS_BLOCK_TABLE_SIZE);

  return block_table;
}

static void
wfs_free_block_table(uint16_t *block_table)
{
  free(block_table);
}

static void
wfs_save_block_table(int fd, uint16_t *block_table)
{
  lseek(fd, WFS_BLOCK_TABLE_START, SEEK_SET);
  write(fd, block_table, WFS_BLOCK_TABLE_SIZE);

  wfs_free_block_table(block_table);
}

/*
 * Implementation of commands.
 */
static int
wfs_create(const char *filename, char **argv)
{
  int i;
  int fd;
  struct stat buf;

  if (stat(filename, &buf) == 0)
    {
      fprintf(stderr, "error: file '%s' already exists\n", filename);
      return -1;
    }

  fd = open(filename, O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
  if (fd < 0)
    {
      fprintf(stderr, "error: could not open file '%s': %s\n",
              filename, strerror(errno));
      return -1;
    }

  /* Clear all file entries */
  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t empty_entry;

      memset(&empty_entry, 0, sizeof(wfs_file_entry_t));
      write(fd, &empty_entry, sizeof(wfs_file_entry_t));
    }

  /* Give file a specified length */
  ftruncate(fd, wfs_get_size());

  close(fd);

  return 0;
}

static int
wfs_list(const char *filename, char **argv)
{
  int i;
  int fd;
  int valid = 0;

  fd = wfs_open_image(filename);
  if (fd < 0)
    return fd;

  printf("%-58s\t%-s\n\n", "Filename", "Size");

  for (i = 0; i < WFS_N_FILES; i++)
    {
      wfs_file_entry_t entry;
      read(fd, &entry, sizeof(wfs_file_entry_t));

      if (!wfs_file_entry_is_empty(&entry))
        {
          if (wfs_file_entry_is_directory(&entry))
            printf("%-58s\tDIR\n", entry.filename);
          else
            printf("%-58s\t%d\n", entry.filename,
                   wfs_file_entry_get_size(&entry));
          valid++;
        }
    }

  close(fd);

  if (valid == 0)
    printf("\tNo files or directories\n");

  return 0;
}

static int
wfs_adddir(const char *filename, char **argv)
{
  int fd;
  const char *directory = argv[0];
  wfs_file_entry_t entry;

  if (strlen(directory) > 58)
    {
      fprintf(stderr, "error: name '%s' too long\n", directory);
      return -1;
    }

  fd = wfs_open_image(filename);
  if (fd < 0)
    return fd;

  /* Check if entry already exists */
  if (wfs_entry_exists(fd, directory))
    {
      fprintf(stderr, "error: entry '%s' already exists\n", directory);
      close(fd);
      return -1;
    }

  /* Set up new directory entry */
  strncpy(entry.filename, directory, 58);
  /* FIXME: allocate block for this directory. */
  entry.start_block = WFS_BLOCK_EOF;
  entry.size = WFS_SIZE_IS_DIRECTORY;

  if (wfs_add_entry(fd, &entry) < 0)
    {
      close(fd);
      return -1;
    }

  close(fd);

  return 0;
}

static int
wfs_put(const char *filename, char **argv)
{
  int fd;
  int local_fd;
  int offset = 0;
  struct stat buf;
  wfs_file_entry_t entry;
  int block;
  uint16_t *block_table = NULL;
  const char *local = argv[0];
  const char *remote = argv[1];

  /* Check if local file exists */
  if (stat(local, &buf) < 0)
    {
      fprintf(stderr, "error: cannot stat '%s': %s\n", local, strerror(errno));
      return -1;
    }

  local_fd = open(local, O_RDONLY);
  if (local_fd < 0)
    {
      fprintf(stderr, "error: cannot open '%s': %s\n", local, strerror(errno));
      return -1;
    }

  /* Open image */
  fd = wfs_open_image(filename);
  if (fd < 0)
    {
      close(local_fd);
      return fd;
    }

  block_table = wfs_load_block_table(fd);

  /* Make sure there is an entry for the remote file */
  strncpy(entry.filename, remote, 58);
  entry.size = buf.st_size;
  block = -1;

  if (!wfs_entry_exists(fd, remote))
    {
      block = wfs_get_free_block(block_table);
      entry.start_block = block + 1;

      if (block < 0)
        {
          fprintf(stderr, "error: no space left on image.\n");

          wfs_free_block_table(block_table);
          close(local_fd);
          close(fd);

          return -1;
        }
    }

  if (wfs_add_or_find_entry(fd, &entry))
    {
      wfs_free_block_table(block_table);
      close(local_fd);
      close(fd);

      return -1;
    }

  if (block < 0)
    block = entry.start_block - 1;

  /* Transfer file */
  while (offset < buf.st_size)
    {
      lseek(fd, WFS_DATA_START + block * WFS_BLOCK_SIZE, SEEK_SET);

      if (buf.st_size - offset > WFS_BLOCK_SIZE)
        {
          int next_block;
          char buffer[WFS_BLOCK_SIZE];

          /* FIXME: Should check return values of read() and write() */
          read(local_fd, buffer, WFS_BLOCK_SIZE);
          write(fd, buffer, WFS_BLOCK_SIZE);
          offset += WFS_BLOCK_SIZE;

          if (block_table[block] != WFS_BLOCK_FREE
              && block_table[block] != WFS_BLOCK_EOF)
            {
              block = block_table[block] - 1;
            }
          else
            {
              block_table[block] = 0x1234; /* Temporary marker */
              next_block = wfs_get_free_block(block_table);

              if (next_block < 0)
                {
                  fprintf(stderr,
                          "error: no space left on image, transfer aborted.\n");

                  block_table[block] = WFS_BLOCK_EOF;

                  wfs_save_block_table(fd, block_table);

                  close(local_fd);
                  close(fd);

                  return -1;
                }

              block_table[block] = next_block + 1;
              block = next_block;
            }
        }
      else
        {
          char buffer[WFS_BLOCK_SIZE];
          int last_block = block;

          /* FIXME: Should check return values of read() and write() */
          read(local_fd, buffer, buf.st_size - offset);
          write(fd, buffer, buf.st_size - offset);
          offset += buf.st_size - offset;

          if (block_table[block] != WFS_BLOCK_FREE
              && block_table[block] != WFS_BLOCK_EOF)
            {
              /* Delete remainder of chain */
              while (block + 1 != WFS_BLOCK_EOF)
                {
                  int clear = block;
                  block = block_table[block] - 1;
                  block_table[clear] = WFS_BLOCK_FREE;
                }
            }

          block_table[last_block] = WFS_BLOCK_EOF;

          break;
        }
    }

  wfs_save_block_table(fd, block_table);

  close(local_fd);
  close(fd);

  return 0;
}

static int
wfs_get(const char *filename, char **argv)
{
  int fd;
  int local_fd;
  int offset = 0;
  wfs_file_entry_t entry;
  int block;
  int ret;
  uint16_t *block_table = NULL;
  const char *remote = argv[0];
  const char *local = argv[1];

  /* Open image */
  fd = wfs_open_image(filename);
  if (fd < 0)
    return fd;

  /* Retrieve file entry */
  ret = wfs_find_entry(fd, remote, &entry);
  if (ret == 0)
    {
      fprintf(stderr, "error: cannot find remote file '%s'\n", remote);
      close(fd);
      return -1;
    }

  /* Prepare local file */
  local_fd = open(local, O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
  if (local_fd < 0)
    {
      fprintf(stderr, "error: cannot open '%s': %s\n", local, strerror(errno));
      close(fd);
      return -1;
    }

  block_table = wfs_load_block_table(fd);
  block = entry.start_block;
  while (offset < wfs_file_entry_get_size(&entry))
    {
      if (block == WFS_BLOCK_FREE || block >= WFS_BLOCK_EOF)
        {
          fprintf(stderr, "error: file system is corrupted, bailing out.\n");

          wfs_free_block_table(block_table);

          close(fd);
          close(local_fd);

          return -1;
        }

      lseek(fd, WFS_DATA_START + (block - 1) * WFS_BLOCK_SIZE, SEEK_SET);

      if (wfs_file_entry_get_size(&entry) - offset > WFS_BLOCK_SIZE)
        {
          char buffer[WFS_BLOCK_SIZE];

          /* FIXME: Should check return values of read() and write() */
          read(fd, buffer, WFS_BLOCK_SIZE);
          write(local_fd, buffer, WFS_BLOCK_SIZE);
          offset += WFS_BLOCK_SIZE;

          block = block_table[block-1];
        }
      else
        {
          char buffer[WFS_BLOCK_SIZE];
          uint32_t read_size;

          read_size = wfs_file_entry_get_size(&entry) - offset;

          /* FIXME: Should check return values of read() and write() */
          read(fd, buffer, read_size);
          write(local_fd, buffer, read_size);
          offset += read_size;

          block = block_table[block-1];
          break;
        }
    }

  wfs_free_block_table(block_table);

  ftruncate(local_fd, wfs_file_entry_get_size(&entry));

  close(local_fd);
  close(fd);

  return 0;
}

static int
wfs_remove(const char *filename, char **argv)
{
  int fd;
  int ret;
  wfs_file_entry_t entry;
  const char *remote = argv[0];

  /* Open image */
  fd = wfs_open_image(filename);
  if (fd < 0)
    return fd;

  /* Retrieve file entry */
  ret = wfs_find_entry(fd, remote, &entry);
  if (ret == 0)
    {
      fprintf(stderr, "error: cannot find remote file '%s'\n", remote);
      close(fd);
      return -1;
    }

  if (!wfs_file_entry_is_directory(&entry))
    {
      int block = entry.start_block;
      uint16_t *block_table = wfs_load_block_table(fd);

      while (block != WFS_BLOCK_EOF)
        {
          int clear = block;
          block = block_table[block-1];
          block_table[clear-1] = WFS_BLOCK_FREE;
        }

      wfs_save_block_table(fd, block_table);
    }

  wfs_clear_entry(fd, &entry);

  close(fd);

  return 0;
}

static int
wfs_fsck(const char *filename, char **argv)
{
  int i;
  int fd;
  uint16_t *block_table;
  int files = 0;
  int directories = 0;
  int status = 0;

  /* Open image */
  fd = wfs_open_image(filename);
  if (fd < 0)
    return fd;

  block_table = wfs_load_block_table(fd);

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      int block;
      int blocks = 0;
      wfs_file_entry_t entry;

      read(fd, &entry, sizeof(wfs_file_entry_t));

      if (entry.start_block != WFS_BLOCK_FREE && entry.filename[0] == 0)
        {
          fprintf(stderr, "File entry %d has non-free start block and no filename\n", i);
          status--;
        }
      if (wfs_file_entry_is_directory(&entry) &&
          wfs_file_entry_get_size (&entry) != 0)
        {
          fprintf(stderr, "File entry %d is a directory but has non-zero size\n", i);
          status--;
        }
      if (!wfs_file_entry_is_directory(&entry) /* FIXME: temporary */
          && entry.start_block == WFS_BLOCK_EOF)
        {
          fprintf(stderr," File entry %d has EOF as start block\n", i);
          status--;
        }

      if (wfs_file_entry_is_empty(&entry)
          || wfs_file_entry_is_directory(&entry))
        {
          if (!wfs_file_entry_is_empty(&entry))
            directories++;
          continue;
        }

      files++;

      /* Check the chain */
      block = entry.start_block;
      blocks = 0;
      while (block != WFS_BLOCK_EOF)
        {
          if (block == WFS_BLOCK_FREE)
            {
              fprintf(stderr, "File %s at %u has a corrupted chain\n", entry.filename, entry.start_block);
              status--;
              break;
            }
          blocks++;
          block = block_table[block-1];
        }

      if (blocks * WFS_BLOCK_SIZE < wfs_file_entry_get_size(&entry))
        {
          fprintf(stderr, "File %s has a broken chain\n", entry.filename);
          status--;
        }
      if (blocks * WFS_BLOCK_SIZE > wfs_file_entry_get_size(&entry) + WFS_BLOCK_SIZE)
        {
          fprintf(stderr, "File %s has a broken chain\n", entry.filename);
          status--;
        }
    }

  wfs_free_block_table(block_table);

  close(fd);

  if (status == 0)
    fprintf(stderr,
            "Check completed. %d files, %d directories. File system is clean.\n",
            files, directories);
  else
    fprintf(stderr,
            "Check completed. %d files, %d directories. %d errors.\n",
            files, directories, status * -1);

  return status;
}

static int
wfs_dumptable(const char *filename, char **argv)
{
  int i;
  int fd;
  uint16_t *block_table;

  /* Open image */
  fd = wfs_open_image(filename);
  if (fd < 0)
    return fd;

  block_table = wfs_load_block_table(fd);

  lseek(fd, 0, SEEK_SET);
  for (i = 0; i < WFS_N_FILES; i++)
    {
      int block;
      wfs_file_entry_t entry;

      read(fd, &entry, sizeof(wfs_file_entry_t));

      if (wfs_file_entry_is_empty(&entry))
        continue;

      printf("%s:\n", entry.filename);
      if (wfs_file_entry_is_directory(&entry))
        {
          printf("Is directory.\n");
          continue;
        }

      block = entry.start_block;
      do
        {
          if (block == WFS_BLOCK_FREE)
            {
              printf("(chain corrupted");
              break;
            }
          printf("%d ", block);
          block = block_table[block-1];
        }
      while (block != WFS_BLOCK_EOF);
      printf("\n\n");
    }

  wfs_free_block_table(block_table);

  close(fd);

  return 0;
}

/*
 * Command registration
 */
struct wfs_command
{
  const char *command;
  const char *arg_description;
  const char *help;
  const int argc;
  int (* func) (const char *filename, char **argv);
}
commands[] =
{
  /* Keep this list alphabetically sorted! */
  { "adddir", "<dirname>", "Add directory", 1, wfs_adddir },
  { "create", "none", "Create file system", 0, wfs_create },
  { "dumptable", "none", "Dump block table", 0, wfs_dumptable },
  { "fsck", "none", "Check file system", 0, wfs_fsck },
  { "get", "<remote> <local>", "Get remote file to local file", 2, wfs_get },
  { "list", "none", "List contents", 0, wfs_list },
  { "put", "<local> <remote>", "Put local file to remote", 2, wfs_put },
  { "remove", "<remote>", "Remove file", 1, wfs_remove },
};

#define N_COMMANDS (sizeof(commands) / sizeof(struct wfs_command))

static int
wfs_command_compare(const void *_key, const void *_member)
{
  const char *key;
  struct wfs_command *member;

  key = (const char *)_key;
  member = (struct wfs_command *)_member;

  return strcmp(key, member->command);
}


static void
print_usage(const char *progname)
{
  int i;

  fprintf(stderr, "usage: %s <wfsfile> command ...\n\n", progname);
  fprintf(stderr, "Where command is one of the following:\n");

  for (i = 0; i < N_COMMANDS; i++)
    {
      fprintf(stderr, "\t%-12s %-16s %s\n",
              commands[i].command, commands[i].arg_description,
              commands[i].help);
    }
}

int
main(int argc, char **argv)
{
  struct wfs_command *command;

  if (argc < 3)
    {
      print_usage(argv[0]);
      return -1;
    }

  command = bsearch(argv[2], commands, N_COMMANDS, sizeof(struct wfs_command),
                    wfs_command_compare);
  if (!command)
    {
      fprintf(stderr, "error: no such command '%s'\n", argv[2]);
      return -1;
    }

  if (argc - 3 < command->argc)
    {
      fprintf(stderr, "error: usage: %s %s\n",
              command->command, command->arg_description);
      return -1;
    }

  return (* command->func) (argv[1], argv + 3);
}
